{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# LSTM修改为二分类，和之前的文本情感分类不一样，之前的是10分类"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "61727\n",
      "ws 加载成功\n",
      "Train Epoch: 0 [0/25000 (0%)]\t损失值: 0.687060\t 准确率: 56.250000\n",
      "Train Epoch: 0 [1280/25000 (5%)]\t损失值: 0.702340\t 准确率: 43.750000\n",
      "Train Epoch: 0 [2560/25000 (10%)]\t损失值: 0.700728\t 准确率: 48.437500\n",
      "Train Epoch: 0 [3840/25000 (15%)]\t损失值: 0.694207\t 准确率: 49.218750\n",
      "Train Epoch: 0 [5120/25000 (20%)]\t损失值: 0.689559\t 准确率: 52.343750\n",
      "Train Epoch: 0 [6400/25000 (26%)]\t损失值: 0.693589\t 准确率: 50.000000\n",
      "Train Epoch: 0 [7680/25000 (31%)]\t损失值: 0.682285\t 准确率: 56.250000\n",
      "Train Epoch: 0 [8960/25000 (36%)]\t损失值: 0.672884\t 准确率: 59.375000\n",
      "Train Epoch: 0 [10240/25000 (41%)]\t损失值: 0.672151\t 准确率: 63.281250\n",
      "Train Epoch: 0 [11520/25000 (46%)]\t损失值: 0.667637\t 准确率: 59.375000\n",
      "Train Epoch: 0 [12800/25000 (51%)]\t损失值: 0.645208\t 准确率: 61.718750\n",
      "Train Epoch: 0 [14080/25000 (56%)]\t损失值: 0.624650\t 准确率: 61.718750\n",
      "Train Epoch: 0 [15360/25000 (61%)]\t损失值: 0.641021\t 准确率: 62.500000\n",
      "Train Epoch: 0 [16640/25000 (66%)]\t损失值: 0.626476\t 准确率: 67.968750\n",
      "Train Epoch: 0 [17920/25000 (71%)]\t损失值: 0.618930\t 准确率: 71.093750\n",
      "Train Epoch: 0 [19200/25000 (77%)]\t损失值: 0.569304\t 准确率: 71.093750\n",
      "Train Epoch: 0 [20480/25000 (82%)]\t损失值: 0.590165\t 准确率: 67.968750\n",
      "Train Epoch: 0 [21760/25000 (87%)]\t损失值: 0.533392\t 准确率: 67.187500\n",
      "Train Epoch: 0 [23040/25000 (92%)]\t损失值: 0.583403\t 准确率: 71.875000\n",
      "Train Epoch: 0 [24320/25000 (97%)]\t损失值: 0.646472\t 准确率: 64.062500\n",
      "\n",
      "测试结果: Avg. 损失: 0.0045, 准确率: 17444/25000 (69.78%)\n",
      "\n",
      "Train Epoch: 1 [0/25000 (0%)]\t损失值: 0.523534\t 准确率: 76.562500\n",
      "Train Epoch: 1 [1280/25000 (5%)]\t损失值: 0.473600\t 准确率: 78.906250\n",
      "Train Epoch: 1 [2560/25000 (10%)]\t损失值: 0.497668\t 准确率: 78.125000\n",
      "Train Epoch: 1 [3840/25000 (15%)]\t损失值: 0.491992\t 准确率: 75.000000\n",
      "Train Epoch: 1 [5120/25000 (20%)]\t损失值: 0.495083\t 准确率: 81.250000\n",
      "Train Epoch: 1 [6400/25000 (26%)]\t损失值: 0.485091\t 准确率: 75.781250\n",
      "Train Epoch: 1 [7680/25000 (31%)]\t损失值: 0.501222\t 准确率: 78.906250\n",
      "Train Epoch: 1 [8960/25000 (36%)]\t损失值: 0.494410\t 准确率: 74.218750\n",
      "Train Epoch: 1 [10240/25000 (41%)]\t损失值: 0.427452\t 准确率: 85.156250\n",
      "Train Epoch: 1 [11520/25000 (46%)]\t损失值: 0.524311\t 准确率: 76.562500\n",
      "Train Epoch: 1 [12800/25000 (51%)]\t损失值: 0.476745\t 准确率: 75.000000\n",
      "Train Epoch: 1 [14080/25000 (56%)]\t损失值: 0.476264\t 准确率: 76.562500\n",
      "Train Epoch: 1 [15360/25000 (61%)]\t损失值: 0.598564\t 准确率: 67.187500\n",
      "Train Epoch: 1 [16640/25000 (66%)]\t损失值: 0.484190\t 准确率: 75.000000\n",
      "Train Epoch: 1 [17920/25000 (71%)]\t损失值: 0.415358\t 准确率: 82.812500\n",
      "Train Epoch: 1 [19200/25000 (77%)]\t损失值: 0.504904\t 准确率: 77.343750\n",
      "Train Epoch: 1 [20480/25000 (82%)]\t损失值: 0.431356\t 准确率: 82.812500\n",
      "Train Epoch: 1 [21760/25000 (87%)]\t损失值: 0.486037\t 准确率: 76.562500\n",
      "Train Epoch: 1 [23040/25000 (92%)]\t损失值: 0.482854\t 准确率: 74.218750\n",
      "Train Epoch: 1 [24320/25000 (97%)]\t损失值: 0.529361\t 准确率: 72.656250\n",
      "\n",
      "测试结果: Avg. 损失: 0.0043, 准确率: 18246/25000 (72.98%)\n",
      "\n",
      "Train Epoch: 2 [0/25000 (0%)]\t损失值: 0.367278\t 准确率: 85.156250\n",
      "Train Epoch: 2 [1280/25000 (5%)]\t损失值: 0.326797\t 准确率: 85.156250\n",
      "Train Epoch: 2 [2560/25000 (10%)]\t损失值: 0.384699\t 准确率: 79.687500\n",
      "Train Epoch: 2 [3840/25000 (15%)]\t损失值: 0.332215\t 准确率: 88.281250\n",
      "Train Epoch: 2 [5120/25000 (20%)]\t损失值: 0.392854\t 准确率: 83.593750\n",
      "Train Epoch: 2 [6400/25000 (26%)]\t损失值: 0.403392\t 准确率: 79.687500\n",
      "Train Epoch: 2 [7680/25000 (31%)]\t损失值: 0.401956\t 准确率: 80.468750\n",
      "Train Epoch: 2 [8960/25000 (36%)]\t损失值: 0.393903\t 准确率: 80.468750\n",
      "Train Epoch: 2 [10240/25000 (41%)]\t损失值: 0.329593\t 准确率: 87.500000\n",
      "Train Epoch: 2 [11520/25000 (46%)]\t损失值: 0.429480\t 准确率: 78.906250\n",
      "Train Epoch: 2 [12800/25000 (51%)]\t损失值: 0.359335\t 准确率: 82.812500\n",
      "Train Epoch: 2 [14080/25000 (56%)]\t损失值: 0.349242\t 准确率: 82.031250\n",
      "Train Epoch: 2 [15360/25000 (61%)]\t损失值: 0.348925\t 准确率: 87.500000\n",
      "Train Epoch: 2 [16640/25000 (66%)]\t损失值: 0.423411\t 准确率: 78.906250\n",
      "Train Epoch: 2 [17920/25000 (71%)]\t损失值: 0.441743\t 准确率: 78.906250\n",
      "Train Epoch: 2 [19200/25000 (77%)]\t损失值: 0.410905\t 准确率: 82.031250\n",
      "Train Epoch: 2 [20480/25000 (82%)]\t损失值: 0.356050\t 准确率: 84.375000\n",
      "Train Epoch: 2 [21760/25000 (87%)]\t损失值: 0.383337\t 准确率: 83.593750\n",
      "Train Epoch: 2 [23040/25000 (92%)]\t损失值: 0.331811\t 准确率: 86.718750\n",
      "Train Epoch: 2 [24320/25000 (97%)]\t损失值: 0.437911\t 准确率: 76.562500\n",
      "\n",
      "测试结果: Avg. 损失: 0.0046, 准确率: 18203/25000 (72.81%)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "import os\n",
    "import re\n",
    "import numpy as np\n",
    "import pickle\n",
    "from tqdm import tqdm\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch import optim\n",
    "\n",
    "data_base_path = r\"data\\aclImdb\"\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "\n",
    "#1. 定义英文单词分类\n",
    "def tokenize(text):\n",
    "    # fileters = '!\"#$%&()*+,-./:;<=>?@[\\\\]^_`{|}~\\t\\n'\n",
    "    fileters = ['!', '\"', '#', '$', '%', '&', '\\(', '\\)', '\\*', '\\+', ',', '-', '\\.', '/', ':', ';', '<', '=', '>',\n",
    "                '\\?', '@'\n",
    "        , '\\[', '\\\\', '\\]', '^', '_', '`', '\\{', '\\|', '\\}', '~', '\\t', '\\n', '\\x97', '\\x96', '”', '“', ]\n",
    "    text = re.sub(\"<.*?>\", \" \", text, flags=re.S)\n",
    "    text = re.sub(\"|\".join(fileters), \" \", text, flags=re.S)\n",
    "    return [i.strip() for i in text.split()]\n",
    "\n",
    "\n",
    "# tqdm可以在长循环中添加一个进度提示信息，用户只需要 封装任意的迭代器 tqdm(iterator)，是一个快速、扩展性强的进度条工具库。\n",
    "# 2.文本序列化\n",
    "class Word2Sequence:\n",
    "    # unknown word 未知的单词，表示在测试集中 新单词\n",
    "    UNK_TAG = \"UNK\"\n",
    "    # padding ，填充的意思，填充的单词\n",
    "    PAD_TAG = \"PAD\"\n",
    "\n",
    "    UNK = 0\n",
    "    PAD = 1\n",
    "\n",
    "    def __init__(self):\n",
    "        self.dict = {\n",
    "            self.UNK_TAG: self.UNK,\n",
    "            self.PAD_TAG: self.PAD\n",
    "        }\n",
    "        self.count = {}\n",
    "        self.isFitted = False\n",
    "\n",
    "    def to_index(self, word):\n",
    "        \"\"\"\n",
    "        单词到数字\n",
    "        :param word:\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        assert self.isFitted == True, \"必须先进行fit操作\"\n",
    "        return self.dict.get(word, self.UNK)\n",
    "\n",
    "    def to_word(self, index):\n",
    "        \"\"\"\n",
    "        数字到单词\n",
    "        :param index:\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        assert self.isFitted, \"必须先进行fit操作\"\n",
    "        if index in self.inversed_dict:\n",
    "            return self.inversed_dict[index]\n",
    "        return self.UNK_TAG\n",
    "\n",
    "    def __len__(self):\n",
    "        return self(self.dict)\n",
    "\n",
    "    def fit(self, sentences):\n",
    "        \"\"\"\n",
    "        把单个单词保存到dict中\n",
    "        :param sentences:[[word1,word2,word3],[word1,word3,wordn..],...]\n",
    "        :param min_count: 最小出现的次数\n",
    "        :param max_count: 最大出现的次数\n",
    "        :param max_feature: 总词语的最大数量\n",
    "        :return:\n",
    "        \"\"\"\n",
    "\n",
    "        # count 表示的是单词的词频\n",
    "        self.isFitted = True\n",
    "        # for sentence in sentences:\n",
    "        #     print(sentence)\n",
    "        for a in sentences:\n",
    "            if a not in self.count:\n",
    "                self.count[a] = 0\n",
    "            self.count[a] += 1\n",
    "\n",
    "    def inverse_transform(self, indices):\n",
    "        \"\"\"\n",
    "        实现从数组 转化为文字\n",
    "        :param indices: [1,2,3....]\n",
    "        :return:[word1,word2.....]\n",
    "        \"\"\"\n",
    "        sentence = []\n",
    "        for i in indices:\n",
    "            word = self.to_word(i)\n",
    "            sentence.append(word)\n",
    "        return sentence\n",
    "\n",
    "    def build_vocab(self, min_count=1, max_feature=None):\n",
    "        \"\"\"\n",
    "        构建dict\n",
    "        :param min_count:\n",
    "        :param max_feature:\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        # 比最小的数量大和比最大的数量小的需要\n",
    "        if min_count is not None:\n",
    "            self.count = {k: v for k, v in self.count.items() if v > min_count}\n",
    "        if max_feature is not None:\n",
    "            self.count = {k: v for k, v in self.count.items() if v <= max_feature}\n",
    "        # 限制最大的数量\n",
    "        if isinstance(max_feature, int):\n",
    "            count = sorted(list(self.count.items()), key=lambda x: x[1])\n",
    "            if max_feature is not None and len(count) > max_feature:\n",
    "                self.count = count[-int(max_feature):]\n",
    "            for w, _ in count:\n",
    "                self.dict[w] = len(self.dict)\n",
    "        else:\n",
    "            for w in sorted(self.count.keys()):\n",
    "                self.dict[w] = len(self.dict)\n",
    "        # 准备一个index->word的字典\n",
    "        self.inversed_dict = dict(zip(self.dict.values(), self.dict.keys()))\n",
    "\n",
    "    def transform(self, sentence, max_len=None):\n",
    "        \"\"\"\n",
    "        实现吧句子转化为数组（向量）\n",
    "        :param sentence:\n",
    "        :param max_len:\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        assert self.isFitted, \"必须先进行fit操作\"\n",
    "\n",
    "        #\n",
    "        if max_len is not None:\n",
    "            r = [self.PAD] * max_len\n",
    "        else:\n",
    "            r = [self.PAD] * len(sentence)\n",
    "\n",
    "        # 如果句子的长度比 max_len长，对句子进行裁剪\n",
    "        if max_len is not None and len(sentence) > max_len:\n",
    "            sentence = sentence[:max_len]\n",
    "\n",
    "        for index, word in enumerate(sentence):\n",
    "            r[index] = self.to_index(word)\n",
    "        return np.array(r, dtype=np.int64)\n",
    "\n",
    "\n",
    "def fit_save_word_sequence():\n",
    "    w = Word2Sequence()\n",
    "    train_path = [os.path.join(data_base_path, i) for i in [\"train/neg\", \"train/pos\"]]\n",
    "    total_file_path_list = []\n",
    "    # 添加路径\n",
    "    for i in train_path:\n",
    "        total_file_path_list.extend([os.path.join(i, j) for j in os.listdir(i)])\n",
    "    # 开始进行分词\n",
    "    for cur_path in tqdm(total_file_path_list, ascii=True, desc=\"fitting\"):\n",
    "        s = tokenize(open(cur_path, encoding=\"UTF-8\").read().strip())\n",
    "        # print(len(s))\n",
    "        w.fit(s)\n",
    "    # 进行\n",
    "    w.build_vocab()\n",
    "    # print(ws.dict)\n",
    "    # print(ws.count)\n",
    "    # 对wordSequesnce进行保存\n",
    "    pickle.dump(w, open(\"./model/ws.pkl\", \"wb\"))\n",
    "\n",
    "\n",
    "# 如果文件不存在，则直接进行存储\n",
    "if not os.path.exists(\"./model/ws.pkl\"):\n",
    "    #2. 在dataset中使用wordsequence\n",
    "    fit_save_word_sequence()\n",
    "    print(\"分词保存成功\")\n",
    "\n",
    "ws = pickle.load(open(\"./model/ws.pkl\", \"rb\"))\n",
    "print(len(ws.count))\n",
    "# ws_test = pickle.load(open(\"./model/ws_test.pkl\", \"rb\"))\n",
    "N = len(ws.dict)\n",
    "\n",
    "max_len = 40\n",
    "print(\"ws 加载成功\")\n",
    "\n",
    "\n",
    "#2. 准备dataset\n",
    "class ImdbDataset(Dataset):\n",
    "    def __init__(self, mode):\n",
    "        super(ImdbDataset, self).__init__()\n",
    "        if mode == \"train\":\n",
    "            text_path = [os.path.join(data_base_path, i) for i in [\"train/neg\", \"train/pos\"]]\n",
    "        else:\n",
    "            text_path = [os.path.join(data_base_path, i) for i in [\"test/neg\", \"test/pos\"]]\n",
    "\n",
    "        self.total_file_path_list = []\n",
    "        for i in text_path:\n",
    "            self.total_file_path_list.extend([os.path.join(i, j) for j in os.listdir(i)])\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        cur_path = self.total_file_path_list[idx]\n",
    "        cur_filename = os.path.basename(cur_path)\n",
    "        # 文件名称为1_1.txt 2_1.txt 后面的1表示情绪的等级\n",
    "        label = int(cur_filename.split(\"_\")[-1].split(\".\")[0]) - 1\n",
    "        # if(temp<=4)\n",
    "        if label <= 4:\n",
    "            label = 0\n",
    "        else:\n",
    "            label = 1\n",
    "        # label = 0 if label <= 4 else 0  #处理标题，获取label，转化为从[0-9]\n",
    "        text = tokenize(open(cur_path, encoding=\"UTF-8\").read().strip())  #直接按照空格进行分词\n",
    "        return label, text\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.total_file_path_list)\n",
    "\n",
    "\n",
    "def collate_fn(batch):\n",
    "    \"\"\"\n",
    "    我们在这里对单词进行 分词解析,把字符串解释成 数字\n",
    "    :param batch:\n",
    "    :return:\n",
    "    \"\"\"\n",
    "    #batch是list，其中是一个一个元组，每个元组是dataset中__getitem__的结果\n",
    "    batch = list(zip(*batch))\n",
    "    labels = torch.LongTensor(batch[0])\n",
    "    # print(labels)\n",
    "    texts = [ws.transform(i, max_len=max_len) for i in batch[1]]\n",
    "    # print(texts)\n",
    "    texts = torch.LongTensor(texts)\n",
    "    del batch\n",
    "    return labels, texts\n",
    "\n",
    "\n",
    "class IMDBLstmmodel(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super(IMDBLstmmodel, self).__init__()\n",
    "        self.hidden_size = 64\n",
    "        self.embedding_dim = 200\n",
    "        self.num_layer = 2\n",
    "        self.bidriectional = True\n",
    "        self.bi_num = 2 if self.bidriectional else 1\n",
    "        self.dropout = 0.5\n",
    "        #以上部分为超参数，可以自行修改\n",
    "\n",
    "        self.embedding = nn.Embedding(N, self.embedding_dim, padding_idx=ws.PAD)  #[N,200]\n",
    "\n",
    "        self.lstm = nn.LSTM(self.embedding_dim,\n",
    "                            self.hidden_size,\n",
    "                            self.num_layer,\n",
    "                            bidirectional=True,\n",
    "                            dropout=self.dropout)\n",
    "        #使用两个全连接层，中间使用relu激活函数\n",
    "        self.fc = nn.Linear(self.hidden_size * self.bi_num, 20)\n",
    "        self.fc2 = nn.Linear(20, 2)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.embedding(x)  # torch.Size([128, 20, 200]) [batch_size, max_len,embedding_dim]\n",
    "        x = x.permute(1, 0, 2)  #进行轴交换[max_len, batch_size, 200]\n",
    "        h_0, c_0 = self.init_hidden_state(x.size(1))  #  torch.Size([4, 128, 64]) torch.Size([4, 128, 64])\n",
    "\n",
    "        # print(h_0.shape, c_0.shape)\n",
    "        _, (h_n, c_n) = self.lstm(x, (h_0, c_0))\n",
    "\n",
    "        #只要最后一个lstm单元处理的结果，这里多去的hidden state\n",
    "        out = torch.cat([h_n[-2, :, :], h_n[-1, :, :]], dim=-1)\n",
    "        out = self.fc(out)\n",
    "        out = F.relu(out)\n",
    "        out = self.fc2(out)\n",
    "        out = F.log_softmax(out, dim=-1)\n",
    "        # print(out.shape)\n",
    "        # torch.Size([128, 2])\n",
    "        return out\n",
    "\n",
    "    def init_hidden_state(self, batch_size):\n",
    "        h_0 = torch.rand(self.num_layer * self.bi_num, batch_size, self.hidden_size)\n",
    "\n",
    "        c_0 = torch.rand(self.num_layer * self.bi_num, batch_size, self.hidden_size)\n",
    "        return h_0, c_0\n",
    "\n",
    "\n",
    "def get_dataloader(mode=True):\n",
    "    \"\"\"\n",
    "    获取数据集加载\n",
    "    :param mode:\n",
    "    :return:\n",
    "    \"\"\"\n",
    "    if mode:\n",
    "        # 2. 实例化，准备dataloader\n",
    "        dataset = ImdbDataset(mode=\"train\")\n",
    "        dataloader = DataLoader(dataset=dataset, batch_size=128, shuffle=True, collate_fn=collate_fn)\n",
    "    else:\n",
    "        dataset = ImdbDataset(mode=\"test\")\n",
    "        dataloader = DataLoader(dataset=dataset, batch_size=128, shuffle=True, collate_fn=collate_fn)\n",
    "    return dataloader\n",
    "\n",
    "\n",
    "train_batch_size = 128\n",
    "test_batch_size = 1000\n",
    "imdb_model = IMDBLstmmodel()\n",
    "optimizer = optim.Adam(imdb_model.parameters())\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "\n",
    "def train(epoch):\n",
    "    mode = True\n",
    "    imdb_model.train(mode)\n",
    "    train_dataloader = get_dataloader(mode)\n",
    "    for idx, (target, input) in enumerate(train_dataloader):\n",
    "        optimizer.zero_grad()\n",
    "        output = imdb_model(input)\n",
    "        # print(target)\n",
    "        loss = F.nll_loss(output, target)  #traget需要是[0,9]，不能是[1-10]\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        if idx % 10 == 0:\n",
    "            pred = torch.max(output, dim=-1, keepdim=False)[-1]\n",
    "            acc = pred.eq(target.data).numpy().mean() * 100.\n",
    "            print('Train Epoch: {} [{}/{} ({:.0f}%)]\\t损失值: {:.6f}\\t 准确率: {:.6f}'.format(epoch, idx * len(input),\n",
    "                                                                                        len(train_dataloader.dataset),\n",
    "                                                                                        100. * idx / len(\n",
    "                                                                                            train_dataloader),\n",
    "                                                                                        loss.item(), acc))\n",
    "            #\n",
    "            # torch.save(imdb_model.state_dict(), \"model/lstm_net.pkl\")\n",
    "            # torch.save(optimizer.state_dict(), 'model/lstm_optimizer.pkl')\n",
    "\n",
    "\n",
    "def test():\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    mode = False\n",
    "    imdb_model.eval()\n",
    "    test_dataloader = get_dataloader(mode)\n",
    "    with torch.no_grad():\n",
    "        for idx, (target, input) in enumerate(test_dataloader):\n",
    "            output = imdb_model(input)\n",
    "            # print(target)\n",
    "            test_loss += F.nll_loss(output, target, reduction=\"mean\")\n",
    "            pred = torch.max(output, dim=-1, keepdim=False)[-1]\n",
    "            correct += pred.eq(target.data).sum()\n",
    "            # acc = 100. * pred.eq(target.data).numpy().mean()\n",
    "            # print('idx: {} Test set: Avg. loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)\\n'.format(idx, test_loss, correct,target.size(0), acc))\n",
    "\n",
    "        test_loss /= len(test_dataloader.dataset)\n",
    "        print('\\n测试结果: Avg. 损失: {:.4f}, 准确率: {}/{} ({:.2f}%)\\n'.format(\n",
    "            test_loss,\n",
    "            correct,\n",
    "            len(test_dataloader.dataset),\n",
    "            100. * correct / len(test_dataloader.dataset)))\n",
    "\n",
    "\n",
    "for i in range(3):\n",
    "    train(i)\n",
    "    test()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}